"""
1. 什么是掩码张量:
    掩代表遮掩，码就是我们张量中的数值，它的尺寸不定，里面一般只有1和0的元素，代表位置被遮掩或者不被遮掩，
    至于是0位置被遮掩还是1位置被遮掩可以自定义，因此它的作用就是让另外一个张量中的一些数值被遮掩，
    也可以说被替换, 它的表现形式是一个张量.

2. 掩码张量的作用:
    在transformer中, 掩码张量的主要作用在应用attention时，
    有一些生成的attention张量中的值计算有可能已知了未来信息而得到的，
    未来信息被看到是因为训练时会把整个输出结果都一次性进行Embedding，
    但是理论上解码器的的输出却不是一次就能产生最终结果的，而是一次次通过上一次结果综合得出的，
    因此，未来的信息可能被提前利用. 所以，我们会进行遮掩. 关于解码器的有关知识将在后面的章节中讲解.
"""

import torch
import numpy as np
import matplotlib.pyplot as plt

def subsequent_mask(size):
    """生成向后遮掩的掩码张量, 参数size是掩码张量最后两个维度的大小, 它的最后两维形成一个方阵"""
    # 在函数中, 首先定义掩码张量的形状
    attn_shape = (1, size, size)

    # 然后使用np.ones方法向这个形状中添加1元素,形成上三角阵, 最后为了节约空间,
    # 再使其中的数据类型变为无符号8位整形unit8
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')

    # 最后将numpy类型转化为torch中的tensor, 内部做一个1 - 的操作,
    # 在这个其实是做了一个三角阵的反转, subsequent_mask中的每个元素都会被1减,
    # 如果是0, subsequent_mask中的该位置由0变成1
    # 如果是1, subsequent_mask中的该位置由1变成0
    return torch.from_numpy(1 - subsequent_mask)

n1 = np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=-1)
n2 = np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=0)
n3 = np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=1)
n4 = np.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], k=2)

print(n1)
print(n2)
print(n3)
print(n4)

size = 5
sm = subsequent_mask(size)
print("sm:", sm)
print(subsequent_mask(20)[0])
plt.figure(figsize=(5,5))
plt.imshow(subsequent_mask(20)[0])


